'''
Version: 1.0
Author: hpf
Date: 2021-05-28 01:44:14
LastEditTime: 2021-07-29 13:02:08
LastEditors: hpf
Description: 超分辨算法运行进程
FilePath: /trackrunner/superrunner.py
'''
import importlib
import json
import os
import sys

import yaml
from PIL import Image
'''
@description: 数据集处理接口
@param {*} library_path 运行的算法路径
@param {*} model_path 运行算法的模型路径
@param {*} config_path 运行算法的配置文件路径
@param {*} input_path 输入算法处理的数据集路径
@param {*} output_path 输出结果的文件夹
@return {*}
'''


def run_eval(library_path, model_path, config_path, input_path, output_path):
    print('Starting run eval...lib:{},model:{},config:{},input:{},output:{}'.
          format(library_path, model_path, config_path, input_path,
                 output_path))

    process_path = sys.argv[0]
    sys.argv.clear()
    sys.argv.append(process_path)

    # 导入算法库
    os.chdir(library_path)
    sys.path.append(library_path)
    sr = importlib.import_module("SR.sr")

    # 加载配置文件
    if config_path == "":
        config_path = library_path + "/test_FaceNet.yml"

    with open(config_path, mode='r') as f:
        opt = yaml.load(f)
        if model_path != "":
            opt["pretrain_model_G"] = model_path
            with open(config_path, mode='w') as w_f:
                yaml.dump(opt, w_f)

    # 初始化变量
    psnr = 0
    ssim = 0
    count = 0

    # 初始化超分x算法
    srer = sr.ZJLAB_SR(config_path)

    # 遍历图片
    for root, dirs, files in os.walk(os.path.join(input_path, 'LR')):
        files_num = len(files)
        for file in files:
            imgpath = os.path.join(root, file)
            image_lr = Image.open(imgpath).convert('RGB')
            bi_img, y_c_img, out_img = srer.run(image_lr)
            out_img.save(os.path.join(output_path, file))
            image_gt = Image.open(os.path.join(input_path, 'GT',
                                               file)).convert('RGB')

            # 运行evaluate计算PSNR和SSIM的值
            psnr_sr, ssim_sr = srer.evaluate(image_gt, out_img)
            print("PlatformProcessNum:%d" % (100 * float(count) / files_num))
            psnr += psnr_sr
            ssim += ssim_sr
            count += 1

    dict = {"FSR PSNR": psnr / count, "FSR SSIM": ssim / count}
    print("PlatformSummary:%s" % (json.dumps(dict)))
    print("PlatformProcessNum:100")


if __name__ == '__main__':
    run_eval(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5])